/* Definitions for the function types.

   This file is part of khipu.

   khipu is free software: you can redistribute it and/or modify
   it under the terms of the GNU Lesser General Public License as published by
   the Free Software Foundation; either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.  */

#include <cstdio>
#include "khipu.hpp"

KP_DECLS_BEGIN

result<object> alloc_fct (interpreter *interp, uint32_t flags)
{
  auto eg = KP_TRY (evh_guard::make (interp));
  auto fp = alloch<function> ();
  fp->bcode = fp->vals = NIL;
  fp->env = fp->name = NIL;
  fp->vo_full |= flags;
  fp->min_argc = fp->max_argc = 0;
  interp->alval = fp->as_obj ();
  gc_register (interp, fp);
  return (interp->alval);
}

bool eq_xx (interpreter *interp, object x, object y)
{
  if ((as_varobj(x)->vo_full ^ as_varobj(y)->vo_full) &
      function_base::native_flag)
    return (false);
  else if (as_varobj(x)->flagged_p (function_base::native_flag))
    return (as_native_fct(x)->fct == as_native_fct(y)->fct);
  else
    return (fct_bcode (x) == fct_bcode (y) &&
            fct_vals (x) == fct_vals (y) && fct_env (x) == fct_env (y));
}

int function::max_stack () const
{
  const bvector *bp = as_bvector (this->bcode);
  auto start = bp->data, end = start + bp->nbytes;
  int max_sp = 0, sp = 0;

  while (start < end)
    {
      auto instr = bcode_get (*start++);
      int opc = instr->opcode ();

      if (instr->loadf_p () || opc == OP_BRBOUND || opc == OP_BRBOUNDL ||
          opc == OP_BIND || opc == OP_BINDL)
        ++sp;
      else if (opc == OP_TRYBEGIN || opc == OP_TRYBEGINL || opc == OP_VFRAME)
        sp += 2;
      else if (instr->branch_p ())
        switch (opc)
          {
            case OP_BRNEQ: case OP_BRNEQL:
              --sp;
            case OP_BRN: case OP_BRNL: case OP_BRT: case OP_BRTL:
              --sp;
          }
      else if (opc == OP_POP || opc == OP_IS || opc == OP_CONS ||
          opc == OP_SETAPOP || opc == OP_SETAPOPL || opc == OP_TRYEND)
        --sp;
      else if (instr->nops () > 0)
        {
          int val = instr->getsarg (start +
            instr->argsize () * (instr->nops () - 1));

          if (opc == OP_KWARGS)
            sp += abs (val) - instr->getsarg (start);
          else if (instr->callf_p () || opc == OP_SKIP || opc == OP_SKIPL)
            sp -= val;
        }

      if (sp > max_sp)
        max_sp = sp;

      start += instr->argsize () * instr->nops ();
    }

  return (max_sp + interpreter::frame_size + 1);
}

result<int64_t> write_x (interpreter *interp, stream *strm,
                         object obj, io_info&)
{
  int64_t ret = KP_TRY (strm->write (interp, "#<function ", 11));
  object name = fct_name (obj);

  if (name == NIL)
    { ret += KP_TRY (strm->write (interp, "#:fct", 5)); }
  else
    { ret += KP_TRY (xwrite (interp, strm, name)); }

  char buf[64];
  ret += KP_TRY (strm->write (interp, buf,
                              sprintf (buf, " at %p>", unmask (obj))));

  return (ret);
}

enum
{
  FCT_PACK_BUILTIN = 0,
  FCT_PACK_BCODE = 1
};

result<int64_t> pack_x (interpreter *interp, stream *strm,
                        object obj, pack_info& info)
{
  if (native_fct_p (obj))
    {
      const char *name = str_cdata (symname (fct_name (obj)));
      int idx = builtin_idx (interp, name) - symbol::N_SPECFORMS;
      int64_t ret = KP_TRY (strm->putb (interp, FCT_PACK_BUILTIN));
      ret += KP_TRY (strm->write (interp, &idx));
      return (ret);
    }

  const function *xp = as_fct (obj);
  int ivs[] =
    {
      (int)(xp->vo_full &
        (function::artificial_flag | function::kwargs_flag |
         function::genericfn_flag)),
      xp->max_sp, xp->min_argc, xp->max_argc
    };

  int64_t ret = KP_TRY (strm->putb (interp, FCT_PACK_BCODE));
  ret += KP_TRY (strm->write (interp, ivs, sizeof (ivs)));
  ret += KP_TRY (xpack (interp, strm, xp->bcode, info));
  ret += KP_TRY (xpack (interp, strm, xp->vals, info));
  ret += KP_TRY ((xp->vo_flags & function::genericfn_flag) ?
                 NIL : xpack (interp, strm, xp->env, info));
  return (ret);
}

static result<object> mcache_make (interpreter *);

result<object> unpack_x (interpreter *interp, stream *strm,
                         pack_info& info, bool save)
{
  valref saved_pos (interp, *info.offset);
  int code = KP_TRY (strm->getb (interp));
  bool rv;

  if (code == FCT_PACK_BUILTIN)
    {
      int idx = -1;

      rv = KP_TRY (strm->sread (interp, &idx));
      if (!rv || (interp->retval = builtin_fct (interp, idx)) == UNBOUND)
        return (info.error ("invalid builtin function index read"));

      return (interp->retval);
    }

  int vals[4];
  {
    auto tmp = KP_TRY (strm->read (interp, vals, sizeof (vals)));
    rv = tmp == sizeof (vals);
  }

  if (!rv)
    return (info.error ("invalid function members read"));

  valref fn = KP_TRY (alloc_fct (interp, vals[0]));
  function *fp = as_fct (*fn);

  fp->bcode = KP_TRY (xunpack (interp, strm, info));
  if (!bvector_p (fp->bcode))
    return (info.error ("invalid function code vector"));

  fp->vals = KP_TRY (xunpack (interp, strm, info));
  if (!array_p (fp->vals) && fp->vals != NIL)
    return (info.error ("invalid function values read"));

  fp->env = KP_TRY (xunpack (interp, strm, info));
  if (!array_p (fp->env) && fp->env != NIL)
    return (info.error ("invalid function environment read"));

  fp->vo_flags = vals[0];
  fp->max_sp = vals[1];
  fp->min_argc = vals[2];
  fp->max_argc = vals[3];

  if (fp->vo_flags & function::genericfn_flag)
    {
      valref tab = KP_TRY (alloc_table (interp, NIL, NIL));
      valref cache = KP_TRY (mcache_make (interp));
      fp->env = KP_TRY (alloc_array (interp, 2));

      xaref(fp->env, 0) = *tab;
      xaref(fp->env, 1) = *cache;
    }

  if (save)
    KP_VTRY (info.add_mapping (interp, *info.offset, *fn));

  kp_return (*fn);
}

const char* fct_sname (object obj)
{
  object name = fct_p (obj) ? fct_name (obj) : NIL;
  return (symbol_p (name) ? str_cdata (symname (name)) : "#:fct");
}

// Generic functions and methods.

/* In our multimethods implementation, each generic function is a closure that
 * stores 2 tables in order to compute the applicable method list.
 *
 * The first table is a basic mapping from a types array to a concrete method,
 * and is updated on each 'defmeth' form, or with a call to 'meth-add'.
 *
 * The second table acts as a cache, and is queried on each call to the
 * generic function. If there is a miss, the applicable methods list is
 * computed at runtime, and the cache is updated for the type list of the
 * arguments that were used on the call to the generic function.
 *
 * The cache has a fixed size, and uses hash chaining in order to make it
 * more compact. The load factor is effectively 2x, but we keep count on
 * the occupancy, and randomly evict entries once that limit is reached.
 * Much like with regular tables, we have both a single-threaded and a
 * multi-thread safe implementation that we pick based on the number
 * of running threads. */

static const uint32_t MCACHE_SIZE = 64;
static const uint32_t MCACHE_LIMIT = MCACHE_SIZE * 2;

/* Cache entry layout:
 * sorted methods | next | types... */

static const uint32_t TYPE_IDX = 2;

static result<object>
mcache_make (interpreter *interp)
{
  object ax = KP_TRY (alloc_array (interp, MCACHE_SIZE + 1, fixint (0)));
  array *rv = as_array (ax);
  ++rv->data, --rv->len;
  *(rv->data - 1) = MCACHE_LIMIT;   // Store limit in a hidden entry.
  kp_return (interp->alval);
}

static uint32_t
mcache_hash (interpreter *interp, const array *key)
{
  /* In order for the hash value to be stable, we compute it on the
   * type names of the key array. */
  uint32_t rv = 0;
  for (uint32_t idx = 0; idx < key->len; ++idx)
    {
      object obj = key->data[idx];
      rv = mix_hash (rv, typespec_p (obj) ?
                     hash_S (interp, type_name (obj)) : 1);
    }

  return (rv);
}

static atomic_t*
mcache_limit (array *ap)
{
  return ((atomic_t *)(ap->data - 1));
}

static inline bool
mcache_eq (const array *x, const array *y)
{
  for (uint32_t i = TYPE_IDX; i < x->len; ++i)
    if (x->data[i] != y->data[i - TYPE_IDX])
      return (false);

  return (true);
}

static object
mcache_get_at (object obj, uint32_t idx, const array *key)
{
  // Returns the method list, or the bucket index in case of a miss.
  while (obj != fixint (0))
    {
      array *ap = as_array (obj);
      if (mcache_eq (ap, key))
        return (ap->data[0]);

      obj = ap->data[1];
    }

  return (fixint (idx));
}

static object
mcache_get (interpreter *interp, const array *cache, const array *key)
{
  uint32_t idx = mcache_hash (interp, key) & (cache->len - 1);
  kp_return (mcache_get_at (cache->data[idx], idx, key));
}

static uint32_t
mcache_bucket_len (object bucket)
{
  uint32_t ret = 0;
  for (; bucket != fixint (0); bucket = xaref (bucket, 1), ++ret) ;
  return (ret);
}

static inline void
mcache_shuffle (interpreter *interp, uint32_t *idxp)
{
  for (uint32_t i = 0; i < MCACHE_SIZE; ++i)
    idxp[i] = i;

  for (uint32_t i = 0; i < MCACHE_SIZE; ++i)
    swap (idxp[i], idxp[(i + interp->xrand ()) % MCACHE_SIZE]);
}

static void
mcache_put_lk (interpreter *interp, array *cache, array *key, uint32_t idx)
{
  // Link the new entry in the bucket.
  key->data[1] = cache->data[idx];
  cache->data[idx] = key->as_obj ();

  if (kp_likely (--*mcache_limit(cache)))
    return;   // There's still room - Done.
  else if (key->data[1] != fixint (0))
    { // The new entry is chained to others - Remove them.
      *mcache_limit(cache) += mcache_bucket_len (key->data[1]);
      key->data[1] = fixint (0);
      return;
    }

  uint32_t ixs[MCACHE_SIZE];
  mcache_shuffle (interp, ixs);

  for (uint32_t i = 0; i < MCACHE_SIZE; ++i)
    {
      uint32_t i2 = ixs[i];
      if (idx == i2)
        continue;   // Don't evict this entry.

      object prev = cache->data[i2];

      if (prev != fixint (0))
        {
          *mcache_limit(cache) += mcache_bucket_len (prev);
          cache->data[i2] = fixint (0);
          break;
        }
    }
}

static result<void>
mcache_clr_lk (interpreter *interp, array *vec)
{
  array *cache = as_array (vec->data[1]);
  memset (cache->data, 0, MCACHE_SIZE * sizeof (*cache->data));
  *mcache_limit(cache) = MCACHE_LIMIT;
  return (0);
}

static void
mcache_put_mt (interpreter *interp, array *cache, array *key, uint32_t idx)
{
  while (true)
    {
      object prev = cache->data[idx];
      key->data[1] = prev;

      if (!fixint_p (mcache_get_at (prev, idx, key)))
        return;   // This very entry was added in the interim.
      else if (atomic_cas_bool ((atomic_t *)&cache->data[idx],
                                prev, key->as_obj ()))
        break;

      atomic_spin_nop ();
    }

  if (atomic_add (mcache_limit (cache), -1) > 0)
    return;

  uint32_t ixs[MCACHE_SIZE];
  mcache_shuffle (interp, ixs);

  for (uint32_t i = 0; i < MCACHE_SIZE; ++i)
    {
      uint32_t i2 = ixs[i];
      if (idx == i2)
        continue;   // Don't evict our own entry.

      object prev = cache->data[i2];
      if (prev != fixint (0) &&
          atomic_cas_bool ((atomic_t *)&cache->data[i2], prev, fixint (0)))
        {
          atomic_add (mcache_limit (cache), mcache_bucket_len (prev));
          break;
        }
    }
}

static result<void>
mcache_clr_mt (interpreter *interp, array *vec)
{
  object nval = KP_TRY (mcache_make (interp));
  vec->data[1] = nval;
  return (gc_wbarrier (interp, vec->as_obj (), nval));
}

struct meth_comparator : public comparator
{
  meth_comparator (interpreter *ip) : comparator (ip)
    {
    }

  result<bool> operator() (object x, object y)
    {
      const array *a = as_array (x), *b = as_array (y);
      for (uint32_t i = 1; i < a->len; ++i)
        if (a->data[i] > b->data[i])
          return (false);

      return (true);
    }
};

static bool
types_apply_p (interpreter *interp, object types, object key, object *dp)
{
  for (uint32_t i = 0; i < len_a (types); ++i)
    {
      object obj = xaref (types, i);
      if (obj == NIL)
        // No type. Distance is set to the maximum.
        *dp++ = fixint (FIXINT_MAX);
      else
        {
          int dist = subtype_p (obj, xaref (key, i));
          if (!dist)
            return (false);

          *dp++ = fixint (dist);
        }
    }

  return (true);
}

static result<object>
methods_applicable (interpreter *interp, object mtable,
                    object types, uint32_t nargs)
{
  tmp_allocator ta { interp };
  table::iterator it { interp, mtable };
  uint32_t nmax = as_table(mtable)->capacity ();
  array *dist = (array *)ta.alloc (nmax * sizeof (*dist));
  object *dp = (object *)ta.alloc (nmax * (nargs + 2) * sizeof (*dp));
  uint32_t ix = 0;

  for (object *tp = dp + nmax; it.valid (); ++it)
    {
      if (!types_apply_p (interp, types, it.key (), tp + 1))
        continue;

      *tp = it.val ();
      dist[ix].data = tp;
      tp += (dist[ix].len = nargs + 1);
      dp[ix] = dist[ix].as_obj ();
      ++ix;
    }

  local_varobj<array> rv;
  rv.local_init (dp, ix);
  valref tmp (interp, rv.as_obj ());

  meth_comparator cmp { interp };
  deref (nsort_a (interp, *tmp, cmp));

  object ret = KP_TRY (alloc_cons (interp, ix));
  for (uint32_t i = 0; i < ix; ++i, ret = xcdr (ret))
    xcar(ret) = xaref (rv.data[i], 0);

  kp_return (interp->alval);
}

static inline bool
meth_vec_p (object vec)
{
  return (array_p (vec) && len_a (vec) == 2 &&
          table_p (xaref (vec, 0)) && array_p (xaref (vec, 1)) &&
          len_a (xaref (vec, 1)) == MCACHE_SIZE);
}

static result<object>
meth_compute (interpreter *interp, object vec, const array *types)
{
  array *cache = as_array (xaref (vec, 1));
  object meths = mcache_get (interp, cache, types);

  if (!fixint_p (meths))
    // Got a cache hit.
    kp_return (meths);

  // The lookup came up empty. Compute the applicable methods.
  int bucket = as_int (meths);
  meths = KP_TRY (methods_applicable (interp, xaref (vec, 0),
                                      types->as_obj (), types->len));
  if (meths == NIL)
    return (meths);

  interp->aux = meths;
  KP_VTRY (alloc_array (interp, TYPE_IDX + types->len, fixint (0)));

  array *ckey = as_array (interp->alval);
  ckey->data[0] = meths;
  copy_objs (&ckey->data[TYPE_IDX], types->data, types->len);
  (singlethr_p () ? mcache_put_lk : mcache_put_mt)
    (interp, cache, ckey, bucket);

  kp_return (meths);
}

static result<object>
meth_compute (interpreter *interp, object vec, object *argv, int argc)
{
  tmp_allocator ta { interp };
  object *tp = (object *)ta.alloc (argc * sizeof (*tp));
  for (int i = 0; i < argc; ++i)
    tp[i] = type (argv[i]);

  local_varobj<array> types;
  types.local_init (tp, argc);
  return (meth_compute (interp, vec, &types));
}

static inline uintptr_t meth_idx ()
{
  return (symtlidx (symbol::meth_curr) - 1);
}

static result<object>
meth_call (interpreter *interp, bool raise, object vec, object *argv, int argc)
{
  object meths = KP_TRY (meth_compute (interp, vec, argv, argc));
  if (meths == NIL)
    {
      if (raise)
        return (interp->raise ("dispatch-error", "no applicable method found"));

      return (UNBOUND);
    }

  // Install the method list and call the first one.
  interp->tl_syms[meth_idx ()] = meths;
  interp->stkend -= interpreter::frame_size;
  *(interp->stkend - argc - 1) = xcar (meths);
  return (call_n (interp, argc));
}

static object
genericfn_env (object fn)
{
  return (xcar (xaref (fct_env (fn), 0)));
}

static inline bool
genericfn_p (object fn)
{
  return (fct_p (fn) && as_fct(fn)->flagged_p (function::genericfn_flag));
}

static inline bool
meth_types_p (object ax, uint32_t argc)
{
  if (!array_p (ax) || len_a (ax) != argc)
    return (false);

  for (uint32_t i = 0; i < argc; ++i)
    if (xaref (ax, i) != NIL && !typespec_p (xaref (ax, i)))
      return (false);

  return (true);
}

static inline bool
meth_compatible_p (const function *gf, const function *meth)
{
  /* A generic function and a method are congruent iff:
   * - The number of required parameters is equal
   * - The generic function has at least the same number of
   *   optional parameters. */
  return (gf->min_argc == meth->min_argc &&
          (uint32_t)gf->max_argc >= (uint32_t)meth->max_argc);
}

static inline result<void>
mcache_clr (interpreter *interp, array *vec, bool mt)
{
  if (*mcache_limit(as_array (vec->data[1])) != MCACHE_LIMIT)
    return ((mt ? mcache_clr_mt : mcache_clr_lk) (interp, vec));
  return (0);
}

static result<void>
meth_add (interpreter *interp, object gfn, object types, object meth)
{
  if (!genericfn_p (gfn))
    return (interp->raise ("arg-error",
                           "second argument must be a generic function"));
  else if (!meth_types_p (types, as_fct(gfn)->min_argc))
    return (interp->raise ("arg-error",
                           "third argument must be a valid type array"));
  else if (!fct_p (meth) || !meth_compatible_p (as_fct (gfn), as_fct (meth)))
    return (interp->raise ("arg-error",
                           "fourth argument must be a compatible function"));

  bool mt = !singlethr_p ();
  array *vec = as_array (genericfn_env (gfn));
  KP_VTRY (table_put (interp, vec->data[0], types, meth, mt));
  return (mcache_clr (interp, vec, mt));
}

static result<object>
meth_del (interpreter *interp, object gfn, object types)
{
  if (!genericfn_p (gfn))
    return (interp->raise ("arg-error",
                           "second argument must be a generic function"));

  bool mt = !singlethr_p ();
  array *vec = as_array (genericfn_env (gfn));
  object rv = KP_TRY (table_pop (interp, vec->data[0], types, NIL, mt));
  KP_VTRY (mcache_clr (interp, vec, mt));
  return (rv == NIL ? rv : symbol::t);
}

static object
meth_test (interpreter *interp, object prev, object nval)
{
  return (genericfn_p (prev) && fct_p (nval) &&
          meth_compatible_p (as_fct (prev), as_fct (nval)) ?
          symbol::t : NIL);
}

result<bool> method_call (interpreter *interp, uint32_t argc)
{
  call_guard g { interp, argc + 1};
  object gfn = *(interp->stkend - argc - 1);

  if (!genericfn_p (gfn))
    return (false);

  object vec = genericfn_env (gfn);
  dbinding curr_meth { interp };

  KP_VTRY (curr_meth.init (meth_idx () + 1, NIL),
           interp->push_frame (NIL, argc, 0),
           meth_call (interp, false, vec, interp->stkend -
                      argc - interpreter::frame_size, argc));
  return (interp->retval != UNBOUND);
}

result<object> p_meth_ctl (interpreter *interp, object *argv, int argc)
{
  if (!fixint_p (*argv))
    return (interp->raise ("type-error", "first argument must be an integer"));

  switch (as_int (*argv))
    {
      case 0:
        if (argc != 3)
          return (interp->raise_nargs (3, 3, argc));
        kp_return (meth_test (interp, argv[1], argv[2]));

      case 1:
        return (mcache_make (interp));

      case 2:
        // (%meth-ctl 2 name method generic-fct)
        if (argc != 4 || !fct_p (argv[2]) || !fct_p (argv[3]) ||
            !array_p (fct_env (argv[3])) ||
            !meth_vec_p (genericfn_env (argv[3])))
          kp_return (NIL);
        else if (symbol_p (argv[1]))
          as_fct(argv[3])->name = argv[1];

        as_fct(argv[3])->set_flag (function::genericfn_flag |
                                   (as_fct(argv[2])->vo_full &
                                    function::kwargs_flag));
        kp_return (argv[3]);

      case 3:
        // (%meth-ctl 3 generic-fct args...)
        interp->cur_frame = as_int (interp->stack[interp->cur_frame - 4]);
        return (meth_call (interp, true, argv[1], argv + 2, argc - 2));

      case 4:
        // (%meth-ctl 4 generic-fct types-array)
        if (argc != 3)
          return (interp->raise_nargs (3, 3, argc));
        else if (!genericfn_p (argv[1]))
          return (interp->raise ("arg-error", "second argument must be a "
                                              "generic function"));
        else if (!meth_types_p (argv[2], as_fct(argv[1])->min_argc))
          return (interp->raise ("arg-error", "third argument must be a "
                                              "compatible types array"));

        return (meth_compute (interp, genericfn_env (argv[1]),
                              as_array (argv[2])));

      case 5:
        // (%meth-ctl 5 generic-fct types-array method)
        if (argc != 4)
          return (interp->raise_nargs (4, 4, argc));

        KP_VTRY (meth_add (interp, argv[1], argv[2], argv[3]));
        kp_return (argv[3]);

      case 6:
        // (%meth-ctl 6 generic-fct types-array)
        if (argc != 3)
          return (interp->raise_nargs (3, 3, argc));

        interp->retval = KP_TRY (meth_del (interp, argv[1], argv[2]));
        return (interp->retval);

      default:
        return (interp->raise ("arg-error", "invalid code specified"));
    }
}

static result<void>
kwargs_set (interpreter *interp, object *argv, int& argc, kwpair *px)
{
  const char *name = px[0].name;
  size_t len = strlen (name);
  for (int i = 0; i < argc; i += 2)
    {
      if (!keyword_p (argv[i]))
        return (interp->raise ("arg-error",
                               KP_SPRINTF (interp,
                                           "expected a keyword, got: %Q",
                                           argv[i])));

      const string *nm = as_str (symname (argv[i]));
      if (nm->nbytes == len && memcmp (nm->data, name, len) == 0)
        {
          *px[1].obj = argv[i + 1];
          swap (argv[i + 0], argv[argc - 2]);
          swap (argv[i + 1], argv[argc - 1]);
          argc -= 2;
          break;
        }
    }

  return (0);
}

result<void> kwargs_parse_pairs (interpreter *interp, object *argv,
                                 int argc, kwpair *pairs, int npair)
{
  for (int i = 0; i < npair; i += 2)
    KP_VTRY (kwargs_set (interp, argv, argc, pairs + i));
  return (0);
}

KP_DECLS_END
